package com.autonavi.collapse.task.db;

import com.autonavi.collapse.task.TaskService;
import com.autonavi.collapse.task.annotation.TaskComponent;
import com.autonavi.collapse.task.dispatch.TaskNode;
import com.autonavi.collapse.task.model.*;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.Timestamp;
import java.util.*;

@TaskComponent
public class DbTaskService implements TaskService {

    private static final Logger LOGGER = LoggerFactory.getLogger(DbTaskService.class);
    private static final String SQL_SELECT_MY_TASK = "SELECT id, type, owner, configs, begin_time, end_time, create_time, status FROM collapse_task WHERE owner=? AND status BETWEEN 0 AND 2 ORDER BY id LIMIT ?";
    private static final String SQL_SELECT_MY_TASK_TYPE = "SELECT id, type, owner, configs, begin_time, end_time, create_time, status FROM collapse_task WHERE owner=? AND type=? AND status BETWEEN 0 AND 2 ORDER BY id LIMIT ?";
    private static final String SQL_SELECT_TASKS = "SELECT id FROM collapse_task WHERE status IN (0,3) AND group_name=? ORDER BY priority,id LIMIT ?";
    private static final String SQL_SELECT_TASKS_TYPE = "SELECT id FROM collapse_task WHERE status IN (0,3) AND group_name=? AND type=? ORDER BY priority,id LIMIT ?";
    private static final String SQL_UPDATE_TASK_STATUS = "UPDATE collapse_task SET status=? WHERE id=? AND owner=?";
    private static final String SQL_UPDATE_TASK_RESULT = "UPDATE collapse_task SET status=?,next_time=?,priority=priority+1 WHERE id=? AND owner=?";
    private static final String SQL_UPDATE_TASK_OWNER = "UPDATE collapse_task SET status=?,owner=? WHERE id=?";
    private static final String SQL_INSERT_TASK = "INSERT collapse_task(group_name, type, configs, begin_time, end_time) VALUES ";
    private static final String SQL_SELECT_TEMPLATE = "SELECT id, name, type, configs, progress, delay, next_time FROM collapse_task_template WHERE group_name=? AND next_time<NOW() AND status>0 LIMIT ?";
    private static final String SQL_UPDATE_TEMPLATE = "UPDATE collapse_task_template SET next_time=?,progress=? WHERE id=?";
    private static final String SQL_SELECT_ALIVE_NODES = "SELECT id, name, free_size FROM collapse_task_node WHERE timeout>NOW() AND status>0 AND group_name=?";
    private static final String SQL_SELECT_DEAD_NODES = "SELECT id, name FROM collapse_task_node WHERE timeout<NOW() AND group_name=?";
    private static final String SQL_HEARTBEAT = "UPDATE collapse_task_node SET free_size=?,timeout=date_add(now(),INTERVAL ? second),group_name=? WHERE name=?";
    private static final String SQL_INIT_NODE = "INSERT collapse_task_node(group_name,name,free_size) VALUES(?,?,?)";
    private static final String SQL_SELECT_TASK_RELATIONS = "SELECT id, task_type, node_name, status FROM collapse_task_relation WHERE group_name=? AND status>0 ORDER BY id LIMIT ?";
    private static final String SQL_SELECT_TASK_TYPE = "SELECT type FROM collapse_task_template WHERE group_name=? AND status=1 GROUP BY type";

    private DataSource dataSource;
    private ObjectMapper objectMapper = new ObjectMapper();
    private String group;

    @Override
    public void init(String group, Object dataCenter) {
        this.group = group;
        if (dataCenter instanceof DataSource) {
            dataSource = (DataSource) dataCenter;
        }
    }

    @Override
    public List<TaskConfig> findMyTask(String myName, int size) {
        List<TaskConfig> result = new ArrayList<>();
        try (Connection connection = dataSource.getConnection();
             PreparedStatement preparedStatement = connection.prepareStatement(SQL_SELECT_MY_TASK)){
            preparedStatement.setString(1, myName);
            preparedStatement.setInt(2, size);
            try (ResultSet resultSet = preparedStatement.executeQuery()) {
                while (resultSet.next()) {
                    Map<String, ConfigValue> map = new HashMap<>();
                    String configs = resultSet.getString(4);
                    if (StringUtils.isNotBlank(configs)) {
                        map.putAll(objectMapper.readValue(configs, new TypeReference<Map<String, ConfigValue>>() {
                        }));
                    }
                    TaskConfig taskConfig = new TaskConfig(map);
                    taskConfig.setId(resultSet.getString(1));
                    taskConfig.setType(resultSet.getString(2));
                    taskConfig.setOwner(resultSet.getString(3));
                    taskConfig.setBeginTime(resultSet.getString(5));
                    taskConfig.setEndTime(resultSet.getString(6));
                    taskConfig.setCreateTime(resultSet.getTimestamp(7));
                    taskConfig.setStatus(resultSet.getInt(8));
                    result.add(taskConfig);
                }
            }
        } catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
        return result;
    }

    @Override
    public List<TaskConfig> findMyTask(String myName, String type, int size) {
        List<TaskConfig> result = new ArrayList<>();
        try (Connection connection = dataSource.getConnection();
             PreparedStatement preparedStatement = connection.prepareStatement(SQL_SELECT_MY_TASK_TYPE)) {
            preparedStatement.setString(1, myName);
            preparedStatement.setString(2, type);
            preparedStatement.setInt(3, size);
            try (ResultSet resultSet = preparedStatement.executeQuery()){
                while (resultSet.next()) {
                    Map<String, ConfigValue> map = new HashMap<>();
                    String configs = resultSet.getString(4);
                    if (StringUtils.isNotBlank(configs)) {
                        map.putAll(objectMapper.readValue(configs, new TypeReference<Map<String, ConfigValue>>() {
                        }));
                    }
                    TaskConfig taskConfig = new TaskConfig(map);
                    taskConfig.setId(resultSet.getString(1));
                    taskConfig.setType(resultSet.getString(2));
                    taskConfig.setOwner(resultSet.getString(3));
                    taskConfig.setBeginTime(resultSet.getString(5));
                    taskConfig.setEndTime(resultSet.getString(6));
                    taskConfig.setCreateTime(resultSet.getTimestamp(7));
                    taskConfig.setStatus(resultSet.getInt(8));
                    result.add(taskConfig);
                }
            }
        }catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
        return null;
    }

    @Override
    public boolean setDispatched(String id, String myName) {
        try (Connection connection = dataSource.getConnection();
            PreparedStatement preparedStatement = connection.prepareStatement(SQL_UPDATE_TASK_STATUS)){
            preparedStatement.setInt(1, TaskConfig.DISPATCHED);
            preparedStatement.setString(2, id);
            preparedStatement.setString(3, myName);
            preparedStatement.execute();
            return preparedStatement.getUpdateCount() > 0;
        }catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
        return false;
    }

    @Override
    public boolean setRunning(String id, String myName) {
        try (Connection connection = dataSource.getConnection();
            PreparedStatement preparedStatement = connection.prepareStatement(SQL_UPDATE_TASK_STATUS)){
            preparedStatement.setInt(1, TaskConfig.RUNNING);
            preparedStatement.setString(2, id);
            preparedStatement.setString(3, myName);
            preparedStatement.execute();
            return preparedStatement.getUpdateCount() > 0;
        }catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
        return false;
    }

    @Override
    public boolean notify(TaskResult taskResult, String myName) {
        int status = -1;
        if (taskResult != null) {
            switch (taskResult.getStatus()) {
                case SUCCESS:
                    status = TaskConfig.FINISH;
                    break;
                case ERROR:
                    status = TaskConfig.ERROR;
                    break;
                case INTERRUPTED:
                    break;
                case FAIL:
                case UNSUPPORTED:
                default:
                    status = TaskConfig.FAIL;
                    break;
            }
        }
        if (status < 0) {
            return false;
        }
        TaskConfig task = taskResult.getTask();
        try (Connection connection = dataSource.getConnection();
            PreparedStatement preparedStatement = connection.prepareStatement(SQL_UPDATE_TASK_RESULT)){
            preparedStatement.setInt(1, status);
            preparedStatement.setTimestamp(2, new Timestamp(System.currentTimeMillis()));
            preparedStatement.setString(3, task.getId());
            preparedStatement.setString(4, myName);
            preparedStatement.execute();
            return preparedStatement.getUpdateCount() > 0;
        }catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
        return false;
    }

    @Override
    public int addTask(List<TaskConfig> taskConfigs) {
        if (taskConfigs == null || taskConfigs.isEmpty()) {
            return 0;
        }
        List<TaskConfig> temp = new ArrayList<>();
        int result = 0;
        for (TaskConfig taskConfig : taskConfigs) {
            temp.add(taskConfig);
            if (temp.size() >= 100) {
                if (addBatch(temp)) {
                    result += temp.size();
                }
                temp = new ArrayList<>();
            }
        }
        if (!temp.isEmpty() && addBatch(temp)) {
            result += temp.size();
        }
        return result;
    }

    private boolean addBatch(List<TaskConfig> temp) {
        List<String> list = new ArrayList<>();
        for (int i =0; i <temp.size(); i++) {
            list.add("(?,?,?,?,?)");
        }
        String sql = SQL_INSERT_TASK + StringUtils.join(list, ",");
        try (Connection connection = dataSource.getConnection();
            PreparedStatement preparedStatement = connection.prepareStatement(sql)){
            int i = 0;
            for (TaskConfig taskConfig : temp) {
                preparedStatement.setString(++i, taskConfig.getGroup());
                preparedStatement.setString(++i, taskConfig.getType());
                preparedStatement.setString(++i, objectMapper.writeValueAsString(taskConfig.getConfigs().getAll()));
                preparedStatement.setString(++i, taskConfig.getBeginTime());
                preparedStatement.setString(++i, taskConfig.getEndTime());
            }
            preparedStatement.execute();
            return preparedStatement.getUpdateCount() > 0;
        }catch (Exception e){
            LOGGER.error(e.getMessage(), e);
        }
        return false;
    }

    @Override
    public List<TaskTemplate> findTemplates(int size) {
        List<TaskTemplate> result = new ArrayList<>();
        try(Connection connection = dataSource.getConnection();
                 PreparedStatement preparedStatement = connection.prepareStatement(SQL_SELECT_TEMPLATE)){
            preparedStatement.setString(1, group);
            preparedStatement.setInt(2, size);
            try (ResultSet resultSet = preparedStatement.executeQuery()){
                while (resultSet.next()) {
                    Map<String, ConfigValue> map = new HashMap<>();
                    String configs = resultSet.getString(4);
                    if (StringUtils.isNotBlank(configs)) {
                        map.putAll(objectMapper.readValue(configs, new TypeReference<Map<String,ConfigValue>>() {
                        }));
                    }
                    TaskTemplate taskTemplate = new TaskTemplate(map);
                    taskTemplate.setId(resultSet.getLong(1));
                    taskTemplate.setName(resultSet.getString(2));
                    taskTemplate.setGroup(group);
                    taskTemplate.setType(resultSet.getString(3));
                    taskTemplate.setProgress(resultSet.getString(5));
                    taskTemplate.setInterval(resultSet.getInt(6));
                    taskTemplate.setNextTime(resultSet.getTimestamp(7));
                    result.add(taskTemplate);
                }
            }
        }catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
        return result;
    }

    @Override
    public boolean saveTemplateNextTime(TaskTemplate taskTemplate) {
        Date nextTime = taskTemplate.getNextTime();
        if (nextTime == null) {
            return false;
        }
        try (Connection connection = dataSource.getConnection();
             PreparedStatement preparedStatement = connection.prepareStatement(SQL_UPDATE_TEMPLATE)){
            preparedStatement.setTimestamp(1, new Timestamp(nextTime.getTime()));
            preparedStatement.setString(2, taskTemplate.getProgress());
            preparedStatement.setLong(3, taskTemplate.getId());
            preparedStatement.execute();
            return preparedStatement.getUpdateCount() > 0;
        }catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
        return false;
    }

    @Override
    public List<TaskNode> findAvailableNodes() {
        List<TaskNode> result = new ArrayList<>();
        try (Connection connection = dataSource.getConnection();
             PreparedStatement preparedStatement = connection.prepareStatement(SQL_SELECT_ALIVE_NODES)){
            preparedStatement.setString(1, group);
            try (ResultSet resultSet = preparedStatement.executeQuery()){
                while (resultSet.next()){
                    TaskNode taskNode = new TaskNode();
                    taskNode.setId(resultSet.getLong(1));
                    taskNode.setName(resultSet.getString(2));
                    taskNode.setFreeSize(resultSet.getInt(3));
                    result.add(taskNode);
                }
            }
        }catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
        return result;
    }

    @Override
    public List<String> findTaskTypeList() {
        List<String> result = new ArrayList<>();
        try (Connection connection = dataSource.getConnection();
             PreparedStatement preparedStatement = connection.prepareStatement(SQL_SELECT_TASK_TYPE)){
            preparedStatement.setString(1, group);
            try (ResultSet resultSet = preparedStatement.executeQuery()){
                while (resultSet.next()){
                    result.add(resultSet.getString(1));
                }
            }
        }catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
        return result;
    }

    @Override
    public List<TaskNode> findDeadNodes() {
        List<TaskNode> result = new ArrayList<>();
        try (Connection connection = dataSource.getConnection();
             PreparedStatement preparedStatement = connection.prepareStatement(SQL_SELECT_DEAD_NODES)){
            preparedStatement.setString(1, group);
            try (ResultSet resultSet = preparedStatement.executeQuery()){
                while (resultSet.next()){
                    TaskNode taskNode = new TaskNode();
                    taskNode.setId(resultSet.getLong(1));
                    taskNode.setName(resultSet.getString(2));
                    result.add(taskNode);
                }
            }
        }catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
        return result;
    }

    @Override
    public List<TaskConfig> findTasks(int size) {
        return getTaskConfigs(size);
    }

    @Override
    public List<TaskConfig> findTasks(int size, String type) {
        return getTaskConfigs(size);
    }

    private List<TaskConfig> getTaskConfigs(int size) {
        List<TaskConfig> result = new ArrayList<>();
        try (Connection connection = dataSource.getConnection();
             PreparedStatement preparedStatement = connection.prepareStatement(SQL_SELECT_TASKS)){
            preparedStatement.setString(1, group);
            preparedStatement.setInt(2, size);
            try (ResultSet resultSet = preparedStatement.executeQuery()){
                while (resultSet.next()){
                    TaskConfig taskConfig = new TaskConfig();
                    taskConfig.setId(resultSet.getString(1));
                    result.add(taskConfig);
                }
            }
        }catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
        return result;
    }

    @Override
    public boolean dispatch(TaskConfig taskConfig) {
        try (Connection connection = dataSource.getConnection();
            PreparedStatement preparedStatement = connection.prepareStatement(SQL_UPDATE_TASK_OWNER)){
            preparedStatement.setInt(1, taskConfig.getStatus());
            preparedStatement.setString(2, taskConfig.getOwner());
            preparedStatement.setString(3, taskConfig.getId());
            preparedStatement.execute();
            return preparedStatement.getUpdateCount() > 0;
        }catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
        return false;
    }

    @Override
    public boolean heartbeat(String myName, int freeSize, int timeout) {
        try (Connection connection = dataSource.getConnection();
        PreparedStatement preparedStatement = connection.prepareStatement(SQL_HEARTBEAT)){
            preparedStatement.setInt(1, freeSize);
            preparedStatement.setInt(2, timeout);
            preparedStatement.setString(3, group);
            preparedStatement.setString(4, myName);
            preparedStatement.execute();
            if (preparedStatement.getUpdateCount() == 0) {
                try (PreparedStatement insertStatement = connection.prepareStatement(SQL_INIT_NODE)){
                    insertStatement.setString(1, group);
                    insertStatement.setString(2, myName);
                    insertStatement.setInt(3, freeSize);
                    insertStatement.execute();
                }
            }
            return true;
        }catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
        return false;
    }

    @Override
    public List<TaskRelation> FindtaskRelations(int size) {
        List<TaskRelation> result = new ArrayList<>();
        try(Connection connection = dataSource.getConnection();
            PreparedStatement preparedStatement = connection.prepareStatement(SQL_SELECT_TASK_RELATIONS)){
            preparedStatement.setString(1, group);
            preparedStatement.setInt(2, size);
            try (ResultSet resultSet = preparedStatement.executeQuery()){
                while (resultSet.next()) {
                    TaskRelation relation = new TaskRelation();
                    relation.setId(resultSet.getLong(1));
                    relation.setTaskType(resultSet.getString(2));
                    relation.setNodeName(resultSet.getString(3));
                    relation.setStatus(resultSet.getInt(4));
                    result.add(relation);
                }
            }
        }catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
        return result;
    }

    public void setDataSource(DataSource dataSource) {
        this.dataSource = dataSource;
    }
}
